/*
 * test KMS API
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */
#include <stdio.h>
#include "attestation.h"
#include "enclave_proxy.h"
#include "kms.h"
#include "qtsm_lib.h"

unsigned long test_kms_gen_random(sig_params_t params, unsigned int rnd_len)
{
    random_data_t rnd_data;
    unsigned long status;
    unsigned int random_len = rnd_len >> 3; /* bit to byte */
    unsigned char random[random_len];
    unsigned int i;

    memset(random, 0, sizeof(random));
    rnd_data.random = random;
    rnd_data.random_len = &random_len;

    /* call kms_gen_random api */
    status = kms_gen_random_blocking(&params, rnd_len, &rnd_data);
    if (status == KMS_SUCCESS) {
        printf("len of rnd_data is %u\n", *(rnd_data.random_len));
        printf("random data is:\n");
        for (i = 0; i < *(rnd_data.random_len); i++)
            printf("%02x ", rnd_data.random[i]);
        printf("\n");
    }

    return status;
}

unsigned long test_kms_generate_datakey(sig_params_t params, keyid_handle_t handle,
    unsigned int datakey_len)
{
    datakey_t res_datakey = {0};
    unsigned long status;
    unsigned int i;

    unsigned int res_plain_key_len = datakey_len;
    unsigned int res_cipher_key_len = 1000;
    unsigned int res_key_id_len = 40;
    char res_key_id[res_key_id_len];
    unsigned char res_plain_key[res_plain_key_len];
    unsigned char res_cipher_key[res_cipher_key_len];

    /* reset datakey memory */
    memset(res_key_id, 0, sizeof(res_key_id));
    memset(res_plain_key, 0, sizeof(res_plain_key));
    memset(res_cipher_key, 0, sizeof(res_cipher_key_len));
    res_datakey.key_id = res_key_id;
    res_datakey.key_id_len = &res_key_id_len;
    res_datakey.plain_key = res_plain_key;
    res_datakey.plain_key_len = &res_plain_key_len;
    res_datakey.cipher_key = res_cipher_key;
    res_datakey.cipher_key_len = &res_cipher_key_len;

    /* call kms generate datakey */
    status = kms_generate_datakey_blocking(&params, &handle, datakey_len, &res_datakey);
    if (status == KMS_SUCCESS) {
        printf("key_id is:\n");
        for (i = 0; i < *(res_datakey.key_id_len); i++)
            printf("%c", res_datakey.key_id[i]);
        printf("\n");

        printf("len of plain_datakey is %u\n", *(res_datakey.plain_key_len));
        printf("plain_datakey is:\n");
        for (i = 0; i < *(res_datakey.plain_key_len); i++)
            printf("%02x", res_datakey.plain_key[i]);
        printf("\n");

        printf("len of cipher_datakey is %u\n", *(res_datakey.cipher_key_len));
        printf("cipher_datakey is:\n");
        for (i = 0; i < *(res_datakey.cipher_key_len); i++)
            printf("%02x", res_datakey.cipher_key[i]);
        printf("\n");
    }

    return status;
}

unsigned long test_kms_decrypt_data(sig_params_t params, keyid_handle_t handle)
{
    plain_cipher_buff_t data_buffer = {0};
    unsigned long status;
    unsigned int i;

    const unsigned char *data_in = "AgBoAGYUAjIU4BEgrXzJm65Ytz51CeW9oMdi0rrH+Q1"
                                   "1i1siY/oBzLxUKrXj3xY74I5sqUq5LkQtEwaOuiIk3I"
                                   "GvIjU3ZmZjOGI3Ny0yYTBjLTRhNWItYTkxYi02ZTNiO"
                                   "DQyZThiMDAAAAAAN3CbGjBY7Fn4mD9GDFoAVxai8Ck3"
                                   "das2/8WnlhnAuVw="; /* hello world! */
    unsigned int data_in_len = strlen(data_in);
    unsigned int data_out_len = 1000;
    unsigned char data_out[data_out_len];

    data_buffer.data_in = data_in;
    data_buffer.data_in_len = data_in_len;
    data_buffer.data_out = data_out;
    data_buffer.data_out_len = &data_out_len;

    /* call kms_decrypt_data api */
    status = kms_decrypt_data_blocking(&params, &handle, &data_buffer);
    if (status == KMS_SUCCESS) {
        printf("plain_data_len is %u\n", *(data_buffer.data_out_len));
        for (i = 0; i < *(data_buffer.data_out_len); i++)
            printf("%c", data_buffer.data_out[i]);
        printf("\n");
    }

    return status;
}

int main(int argc, char *argv[])
{
    sig_params_t params;
    keyid_handle_t handle;
    char *keyid = "7ffc8b77-xxxxxxxxxxxxxxxx";
    unsigned long status;
    int rc = NO_ERROR;
    unsigned int rnd_len = 256;
    unsigned int datakey_length = 512;
    struct connect_info conn;

    memset(&conn, 0, sizeof(struct connect_info));

    if (argc < 3) {
        printf("qtsm_demo vsock_cid vsock_port\n");
        return -1;
    }

    int vsock_cid = atoi(argv[1]);
    int vsock_port = atoi(argv[2]);

    sig_params_init(&params);
    sig_str_t ak = sig_str("VBBxxxxxxxxxxx");
    sig_str_t sk = sig_str("FvI0lxxxxxxxxxxxxxx");
    sig_str_t host = sig_str("kms.br-iaas-odin1.myhuaweicloud.com");
    sig_str_t uri_prefix = sig_str("/v1.0/06057684d580d5762f3ec002a5fcac18/kms/");

    /* Set sign params */
    params.key = ak;
    params.secret = sk;
    params.host = host;
    params.uri_prefix = uri_prefix;

    /* Set key_id */
    handle.key_id = keyid;
    handle.len = strlen(keyid);

    /* setup proxy */
    rc = setup_proxy(&conn, vsock_cid, vsock_port);
    if (rc != NO_ERROR) {
        printf("setup_proxy error.\n");
        return -1;
    }

    /* test gen random api */
    status = test_kms_gen_random(params, rnd_len);
    if (status != KMS_SUCCESS) {
        printf("kms gen-random api failed.\n");
        return -1;
    }

    /* test gen datakey api */
    status = test_kms_generate_datakey(params, handle, datakey_length);
    if (status != KMS_SUCCESS) {
        printf("kms gen datakey api failed.\n");
        goto err;
    }

    /* test decrypt data api */
    status = test_kms_decrypt_data(params, handle);
    if (status != KMS_SUCCESS) {
        printf("kms gen datakey api failed.\n");
        goto err;
    }

err:
    /* close proxy */
    close_proxy(&conn);

    return 0;
}
